/*
 * Copyright (C) 2020 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import { IOException, ProtocolException } from './Error/error'
import { Headers } from './Headers'
import { MediaType } from './MediaType'
import { ResponseBody } from './ResponseBody'
import { BufferedSource } from './okio/BufferedSource'
import { Source } from './okio/Source'
import { HeadersReader } from './http1/HeadersReader'
import { util } from '@kit.ArkTS'
import { StringToUint8Array } from './utils/StringUtil'

/**
 * Reads a stream of [RFC 2046][rfc_2046] multipart body parts. Callers read parts one-at-a-time
 * until [nextPart] returns null. After calling [nextPart] any preceding parts should not be read.
 *
 * Typical use loops over the parts in sequence:
 *
 * ```
 * const response: Response = await call.execute()
 * const multipartReader = new MultipartReader(response.body!)
 *
 * try {
 *   while (true) {
 *     const part = await multipartReader.nextPart()
 *     if (!part) break
 *     await process(part.headers, part.body)
 *   }
 * } finally {
 *   multipartReader.close()
 * }
 * ```
 *
 * Note that [nextPart] will skip any unprocessed data from the preceding part. If the preceding
 * part is particularly large or if the underlying source is particularly slow, the [nextPart] call
 * may be slow!
 *
 * Closing a part **does not** close this multipart reader; callers must explicitly close this with
 * [close].
 *
 * [rfc_2046]: http://www.ietf.org/rfc/rfc2046.txt
 */
export class MultipartReader {
  readonly source: BufferedSource
  private readonly boundary: string

  /** This delimiter typically precedes the first part. */
  private readonly dashDashBoundary: Uint8Array

  /**
   * This delimiter typically precedes all subsequent parts. It may also precede the first part
   * if the body contains a preamble.
   */
  private readonly crlfDashDashBoundary: Uint8Array

  private partCount = 0
  private closed = false
  private noMoreParts = false
  /** This is only part that's allowed to read from the underlying source. */
  currentPart: PartSource | null = null

  /**
   * Creates a MultipartReader from a ResponseBody.
   * @throws IOException if the Content-Type doesn't have a boundary parameter
   */
  constructor(response: ResponseBody)
  /**
   * Creates a MultipartReader from a BufferedSource and boundary string.
   * @param source The source to read from
   * @param boundary The boundary string
   */
  constructor(source: BufferedSource, boundary: string)
  constructor(sourceOrResponse: BufferedSource | ResponseBody, boundary?: string) {
    if (sourceOrResponse instanceof ResponseBody) {
      const response = sourceOrResponse
      const contentType = response.contentType()
      const boundaryParam = contentType?.parameter("boundary")
      if (!boundaryParam) {
        throw new ProtocolException("expected the Content-Type to have a boundary parameter")
      }
      this.source = response.source().buffer()
      this.boundary = boundaryParam
    } else {
      this.source = sourceOrResponse
      this.boundary = boundary!
    }

    // Build "--" + boundary
    const dashDash = StringToUint8Array("--")
    const boundaryBytes = StringToUint8Array(this.boundary)
    this.dashDashBoundary = new Uint8Array(dashDash.length + boundaryBytes.length)
    this.dashDashBoundary.set(dashDash)
    this.dashDashBoundary.set(boundaryBytes, dashDash.length)

    // Build "\r\n--" + boundary
    const crlf = StringToUint8Array("\r\n")
    this.crlfDashDashBoundary = new Uint8Array(crlf.length + dashDash.length + boundaryBytes.length)
    this.crlfDashDashBoundary.set(crlf)
    this.crlfDashDashBoundary.set(dashDash, crlf.length)
    this.crlfDashDashBoundary.set(boundaryBytes, crlf.length + dashDash.length)
  }

  /**
   * Reads the next part from the multipart stream.
   * @returns The next part, or null if there are no more parts
   * @throws IOException if there's an error reading the stream
   */
  async nextPart(): Promise<Part | null> {
    if (this.closed) {
      throw new Error("closed")
    }
    if (this.noMoreParts) {
      return null
    }

    // Read a boundary, skipping the remainder of the preceding part as necessary.
    if (this.partCount === 0 && await this.rangeEquals(0, this.dashDashBoundary)) {
      // This is the first part. Consume "--" followed by the boundary.
      await this.skip(this.dashDashBoundary.length)
    } else {
      // This is a subsequent part or a preamble. Skip until "\r\n--" followed by the boundary.
      while (true) {
        const toSkip = await this.currentPartBytesRemaining(8192)
        if (toSkip === 0) break
        await this.skip(toSkip)
      }
      await this.skip(this.crlfDashDashBoundary.length)
    }

    // Read either \r\n or --\r\n to determine if there is another part.
    let whitespace = false
    let found = false
    let optionIndex = -1

    while (!found) {
      optionIndex = await this.select(afterBoundaryOptions)
      switch (optionIndex) {
        case 0: {
          // "\r\n": We've found a new part.
          this.partCount++
          found = true
          break
        }
        case 1: {
          // "--": No more parts.
          if (whitespace) {
            throw new ProtocolException("unexpected characters after boundary")
          }
          if (this.partCount === 0) {
            throw new ProtocolException("expected at least 1 part")
          }
          this.noMoreParts = true
          return null
        }
        case 2:
        case 3: {
          // " " or "\t" Ignore whitespace and keep looking.
          whitespace = true
          break
        }
        case -1: {
          throw new ProtocolException("unexpected characters after boundary")
        }
      }
    }

    // There's another part. Parse its headers and return it.
    const headers = await new HeadersReader(this.source).readHeaders()
    const partSource = new PartSource(this)
    this.currentPart = partSource
    return new MultipartReader.Part(headers, partSource.buffer())
  }

  /**
   * Returns a value in [0..maxByteCount] with the number of bytes that can be read from [source] in
   * the current part. If this returns 0 the current part is exhausted; otherwise it has at least
   * one byte left to read.
   */
  async currentPartBytesRemaining(maxResult: number): Promise<number> {
    await this.require(this.crlfDashDashBoundary.length)
    const delimiterIndex = await this.indexOf(this.crlfDashDashBoundary)
    if (delimiterIndex === -1) {
      return Math.min(maxResult, this.source.bufferedSize - this.crlfDashDashBoundary.length + 1)
    } else {
      return Math.min(maxResult, delimiterIndex)
    }
  }

  /**
   * Checks if the buffer starts with the given bytes at the given offset.
   * This peeks at the buffer without consuming data.
   */
  private async rangeEquals(offset: number, bytes: Uint8Array): Promise<boolean> {
    await this.require(offset + bytes.length)
    const buffer = this.source.buffer()
    // Access internal state using type assertion (not ideal but necessary)
    const bufferedSource = buffer as any
    const start = bufferedSource.start + offset
    const end = start + bytes.length
    const uint8array = bufferedSource.uint8array

    if (end > bufferedSource.end) {
      return false
    }

    for (let i = 0; i < bytes.length; i++) {
      if (uint8array[start + i] !== bytes[i]) {
        return false
      }
    }
    return true
  }

  /**
   * Finds the index of the given bytes in the buffer, or -1 if not found.
   */
  private async indexOf(bytes: Uint8Array): Promise<number> {
    await this.require(bytes.length)
    const buffer = this.source.buffer()
    const bufferedSource = buffer as any
    const start = bufferedSource.start
    const end = bufferedSource.end
    const uint8array = bufferedSource.uint8array

    if (end - start < bytes.length) {
      return -1
    }

    for (let i = start; i <= end - bytes.length; i++) {
      let match = true
      for (let j = 0; j < bytes.length; j++) {
        if (uint8array[i + j] !== bytes[j]) {
          match = false
          break
        }
      }
      if (match) {
        return i - start
      }
    }
    return -1
  }

  /**
   * Selects the first matching option from the given options array.
   * Returns the index of the matched option, or -1 if none match.
   * 
   * **Important**: Unlike just checking, this method also consumes the matched bytes,
   * matching the behavior of Okio's Options.select().
   */
  private async select(options: Uint8Array[]): Promise<number> {
    // Find the maximum length option to ensure we have enough data
    const maxLength = Math.max(...options.map(opt => opt.length))
    await this.require(maxLength)
    
    const buffer = this.source.buffer()
    const bufferedSource = buffer as any
    const start = bufferedSource.start
    const end = bufferedSource.end
    const uint8array = bufferedSource.uint8array

    for (let optionIndex = 0; optionIndex < options.length; optionIndex++) {
      const option = options[optionIndex]
      if (end - start < option.length) {
        continue
      }

      let match = true
      for (let i = 0; i < option.length; i++) {
        if (uint8array[start + i] !== option[i]) {
          match = false
          break
        }
      }

      if (match) {
        // Consume the matched bytes (this is what Okio's select does)
        await this.skip(option.length)
        return optionIndex
      }
    }

    return -1
  }

  /**
   * Skips the given number of bytes.
   */
  private async skip(byteCount: number): Promise<void> {
    const sink = new Uint8Array(byteCount)
    const bytesRead = await this.source.read(sink, byteCount)
    if (bytesRead !== byteCount) {
      throw new IOException(`Failed to skip ${byteCount} bytes, only skipped ${bytesRead}`)
    }
  }

  /**
   * Ensures that at least the given number of bytes are available in the buffer.
   * This will read from the underlying source if necessary.
   * 
   * Matches Okio's BufferedSource.require() behavior.
   */
  private async require(byteCount: number): Promise<void> {
    // Check if we already have enough data in the buffer
    if (this.source.bufferedSize >= byteCount) {
      return
    }
    
    // If the source is closed and we don't have enough data, throw
    if (this.source.closed) {
      throw new IOException(`Required ${byteCount} bytes but only ${this.source.bufferedSize} available and source is closed`)
    }
    
    // For BufferedSource, the data should already be buffered or will be buffered on access
    // If we don't have enough data, it means the stream ended prematurely
    throw new IOException(`Required ${byteCount} bytes but only ${this.source.bufferedSize} available`)
  }

  /**
   * Closes the multipart reader and releases resources.
   */
  close(): void {
    if (this.closed) return
    this.closed = true
    this.currentPart = null
    this.source.close()
  }

  /**
   * A single part in a multipart body.
   */
  static Part = class Part {
    readonly headers: Headers
    readonly body: BufferedSource

    constructor(headers: Headers, body: BufferedSource) {
      this.headers = headers
      this.body = body
    }

    /**
     * Closes this part's body.
     */
    close(): void {
      this.body.close()
    }
  }
}

export type Part = InstanceType<typeof MultipartReader.Part>

/**
 * A Source implementation that reads only the current part's data.
 */
class PartSource implements Source {
  private readonly multipartReader: MultipartReader
  closed = false
  private bufferedWrapper: PartBufferedSource | null = null

  constructor(multipartReader: MultipartReader) {
    this.multipartReader = multipartReader
  }

  async read(sink: Uint8Array, byteCount: number): Promise<number> {
    if (byteCount < 0) {
      throw new Error(`byteCount < 0: ${byteCount}`)
    }
    if (this.closed || this.multipartReader.currentPart !== this) {
      throw new Error("closed")
    }

    const limit = await this.multipartReader.currentPartBytesRemaining(byteCount)
    if (limit === 0) {
      return -1 // No more bytes in this part.
    }

    return await this.multipartReader.source.read(sink, limit)
  }

  readBytes(): Promise<Uint8Array> {
    throw new Error("Method not implemented.")
  }

  readString(): Promise<string> {
    throw new Error("Method not implemented.")
  }

  Stream(callback: (chunk: Uint8Array, complete?: boolean, error?: Error) => void): void {
    throw new Error("Method not implemented.")
  }

  close(): void {
    if (this.closed) return
    this.closed = true
    if (this.multipartReader.currentPart === this) {
      this.multipartReader.currentPart = null
    }
  }

  buffer(): BufferedSource {
    if (!this.bufferedWrapper) {
      this.bufferedWrapper = new PartBufferedSource(this)
    }
    return this.bufferedWrapper
  }

  skipAll(): void {
    // Skip all remaining data in this part
    // Implementation would need to read and discard data
  }

  _isGzip: boolean = false
  set isGzip(value: boolean) {
    this._isGzip = value
  }
}

/**
 * A BufferedSource wrapper around PartSource.
 */
class PartBufferedSource extends BufferedSource {
  private readonly partSource: PartSource

  constructor(partSource: PartSource) {
    super(null, 8192)
    this.partSource = partSource
  }

  override async read(sink: Uint8Array, byteCount: number): Promise<number> {
    return await this.partSource.read(sink, byteCount)
  }

  override close(): void {
    this.partSource.close()
    super.close()
  }

  override get bufferedSize(): number {
    // For a wrapped source, we can't know the buffered size without reading
    // Return 0 to indicate unknown, or try to peek at the underlying source
    return 0
  }

  override get closed(): boolean {
    return this.partSource.closed || super.closed
  }

  override get exhausted(): boolean {
    return this.closed
  }
}

/**
 * These options follow the boundary.
 */
const afterBoundaryOptions: Uint8Array[] = [
  StringToUint8Array("\r\n"), // 0.  "\r\n"  More parts.
  StringToUint8Array("--"),   // 1.  "--"    No more parts.
  StringToUint8Array(" "),    // 2.  " "     Optional whitespace. Only used if there are more parts.
  StringToUint8Array("\t")    // 3.  "\t"    Optional whitespace. Only used if there are more parts.
]

